#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include "utils.h"
#include "holger_time.h"
#include "ultrasound_simulator.h"
#include <omp.h>

extern unsigned char * mask;
extern int bscan_w, bscan_h;
extern float * cal_matrix;
extern float bscan_spacing_x, bscan_spacing_y;

#define pos_matrix_a(x,y) (pos_matrix[(y)*4 + (x)])
#define inrange(x,a,b) ((x) >= (a) && (x) < (b))
#define volume_a(x,y,z) (volume[(x) + (y)*volume_w + (z)*volume_w*volume_h])
#define max3(a,b,c) max(a, max(b, c))
#define bscans_queue_a(n, x, y) bscans_queue[n][(x) + (y)*bscan_w]

//#define PRINT
#define PRINT_RES0 3
#define PRINT_RES1 30
#define PRINT_RES2 6000

#define dw_n 8
#define dw_now 4

int volume_w = 450;
int volume_h = 200;
int volume_n = 450;
unsigned char * volume = (unsigned char *) malloc(sizeof(unsigned char)*volume_w*volume_h*volume_n);
float volume_spacing = 0.1;
float3 volume_origo = {90.0f, -20.0f, 280.0f}; // ultrasoundSample5

plane_pts * plane_points_queue = (plane_pts *) malloc(dw_n*sizeof(plane_pts));
plane_eq * bscan_plane_equation_queue = (plane_eq *) malloc(dw_n*sizeof(plane_pts));
float3 * intersections = (float3 *) malloc(sizeof(float3)*2*max3(volume_w, volume_h, volume_n)*max3(volume_w, volume_h, volume_n));
unsigned char * * bscans_queue = (unsigned char * *) malloc(dw_n * bscan_w * bscan_h * sizeof(unsigned char));
float * * pos_matrices_queue = (float * *) malloc(dw_n * sizeof(float)*12);
float * pos_timetags_queue = (float *) malloc(dw_n * sizeof(float));
float * bscan_timetags_queue = (float *) malloc(dw_n * sizeof(float));

void shift_queues() {
	for (int i = 0; i < dw_n-1; i++) {
		bscans_queue[i] = bscans_queue[i+1];
		pos_matrices_queue[i] = pos_matrices_queue[i+1];
		bscan_timetags_queue[i] = bscan_timetags_queue[i+1];
		pos_timetags_queue[i] = pos_timetags_queue[i+1];
		bscan_plane_equation_queue[i] = bscan_plane_equation_queue[i+1];
		plane_points_queue[i] = plane_points_queue[i+1];
	}
}

void wait_for_input() {
	bscans_queue[dw_n-1] = NULL;
	pos_matrices_queue[dw_n-1] = NULL;
	while (bscans_queue[dw_n-1] == NULL || pos_matrices_queue[dw_n-1] == NULL) {
		if (poll_bscan()) get_last_bscan(&bscan_timetags_queue[dw_n-1], &bscans_queue[dw_n-1]);
		if (poll_pos_matrix()) get_last_pos_matrix(&pos_timetags_queue[dw_n-1], &pos_matrices_queue[dw_n-1]);
	}
}

void calibrate_pos_matrix(float * pos_matrix, float * cal_matrix) {
	// Multiply cal_matrix into pos_matrix
	float * new_matrix = (float *) malloc(sizeof(float)*12);
	for (int b = 0; b < 3; b++)
		for (int c = 0; c < 4; c++) {
			float sum = 0;
			for (int k = 0; k < 3; k++) 
				sum += cal_matrix[b*4 + k]*pos_matrix[k*4 + c];	
			new_matrix[b*4 + c] = sum;
		}
	memcpy(pos_matrix, new_matrix, 12*sizeof(float));
	free(new_matrix);
}

void create_plane_points(float * pos_matrix) {
	// Fill plane_points
	float3 temp0 = {0.0f,0.0f,0.0f};
	plane_points_queue[dw_n-1].corner0 = temp0;
	float3 temp1 = {0.0f, bscan_w*bscan_spacing_x, 0.0f};
	plane_points_queue[dw_n-1].cornerx = temp1;
	float3 temp2 = {0.0f, 0.0f, bscan_h*bscan_spacing_y};
	plane_points_queue[dw_n-1].cornery = temp2;

	#ifdef PRINT
	printf("plane_points:\n ");
	for (int i = 0; i < 3; i++) {	
		printf("%5.2f ", plane_points[i].x);
		printf("%5.2f ", plane_points[i].y);
		printf("%5.2f ", plane_points[i].z);
		printf(" \t ");
	}
	printf("\n");
	#endif

	// Transform plane_points
	float3 * foo = (float3 *) &plane_points_queue[dw_n-1];
	float * sums = (float *) malloc(sizeof(float)*3);	
	for (int i = 0; i < 3; i++) {	
		for (int y = 0; y < 3; y++) {
			float sum = 0;
			sum += pos_matrix_a(0,y)*foo[i].x;
			sum += pos_matrix_a(1,y)*foo[i].y;
			sum += pos_matrix_a(2,y)*foo[i].z;
			sum += pos_matrix_a(3,y);
			sums[y] = sum;
		}
		memcpy(&foo[i], sums, 3*sizeof(float));
		foo[i] = sub(foo[i], volume_origo);
	}

	#ifdef PRINT
	printf("Transformed translated translated plane_points:\n ");
	for (int i = 0; i < 3; i++) {	
		printf("%5.2f ", plane_points[i].x);
		printf("%5.2f ", plane_points[i].y);
		printf("%5.2f ", plane_points[i].z);
		printf(" \t ");
	}
	printf("\n");
	#endif
}

void create_plane_eq() {
	// Fill bscan_plane_equation
	float3 a = plane_points_queue[dw_n-1].corner0;
	float3 b = plane_points_queue[dw_n-1].cornerx;
	float3 c = plane_points_queue[dw_n-1].cornery;
	float3 _normal = normalize(cross(sub(a,b), sub(c,a)));

	bscan_plane_equation_queue[dw_n-1].a = _normal.x;
	bscan_plane_equation_queue[dw_n-1].b = _normal.y;
	bscan_plane_equation_queue[dw_n-1].c = _normal.z;
	bscan_plane_equation_queue[dw_n-1].d = -_normal.x*a.x - _normal.y*a.y - _normal.y*a.z;

	#ifdef PRINT
	printf("bscan_plane_equation:\n %f %f %f %f\n", bscan_plane_equation.a, bscan_plane_equation.b, bscan_plane_equation.c, bscan_plane_equation.d);
	#endif
}

int find_intersections(int axis) {
	int intersection_counter = 0;

	float3 Rd = {0, 0, 0};
	if (axis == 0) Rd.x = 1;
	if (axis == 1) Rd.y = 1;
	if (axis == 2) Rd.z = 1;

	int y = 0;
	for (int x = 0; x < volume_w; x++) {
		for (int z = 0; z < volume_n; z++) {
			for (int i = 0; i < 2; i++) {
				float3 Pn = {bscan_plane_equation_queue[2-i].a, bscan_plane_equation_queue[2-i].b, bscan_plane_equation_queue[dw_now-i].c};
				float3 R0 = {x*volume_spacing, y*volume_spacing, z*volume_spacing};
				float Vd = dot(Pn, Rd);
				float V0 = -(dot(Pn, R0) + bscan_plane_equation_queue[dw_now-i].d);
				float t = V0/Vd;

				float3 intersection = add(R0, scale(t, Rd));
				intersections[intersection_counter*2 + i] = intersection;

				#ifdef PRINT
				static int print_counter = 0;
				if (print_counter++ % PRINT_RES2 == 0) {
					printf("R0/t: %.2f %.2f %.2f  %.2f\t", R0.x, R0.y, R0.z, t);
					printf("int.: %.2f %.2f %.2f\n", intersection.x, intersection.y, intersection.z);
				}
				#endif
			}
			// TODO: swap the two intersections if they are not ordered
			intersection_counter++;
		}	
	}

	return intersection_counter;
}

bool project(unsigned char * bilinear, float * dist, float3 voxel_coord, int q_idx) {
	float3 normal = {bscan_plane_equation_queue[q_idx].a, bscan_plane_equation_queue[q_idx].b, bscan_plane_equation_queue[q_idx].c};
	float3 x_vector = normalize(sub(plane_points_queue[q_idx].cornerx, plane_points_queue[q_idx].corner0));
	float3 y_vector = normalize(sub(plane_points_queue[q_idx].cornery, plane_points_queue[q_idx].corner0));

	float dist0 = abs(distance(voxel_coord, bscan_plane_equation_queue[q_idx]));
	float3 p0 = sub(add(voxel_coord, scale(-dist0, normal)), plane_points_queue[q_idx].corner0);
	float px0 = dot(p0, x_vector)/bscan_spacing_x;
	float py0 = dot(p0, y_vector)/bscan_spacing_y;
	float xa = px0-floor(px0);
	float ya = py0-floor(py0);
	int xa0 = (int)px0;
	int ya0 = (int)py0;

	bool valid = false;
	unsigned char bilinear0;

	if (inrange(xa0, 0, bscan_w) && inrange(ya0, 0, bscan_h) && inrange(xa0+1, 0, bscan_w) && inrange(ya0+1, 0, bscan_h))
		if (mask[xa0 + ya0*bscan_w] != 0 && mask[xa0+1 + (ya0+1)*bscan_w] != 0 && mask[xa0+1 + ya0*bscan_w] != 0 && mask[xa0 + (ya0+1)*bscan_w] != 0) {
			bilinear0 = bscans_queue_a(q_idx,xa0,ya0)*(1-xa)*(1-ya) + bscans_queue_a(q_idx,xa0+1,ya0)*xa*(1-ya) + bscans_queue_a(q_idx,xa0,ya0+1)*(1-xa)*ya + bscans_queue_a(q_idx,xa0+1,ya0+1)*xa*ya;
			valid = true;
		}

	*dist = dist0;
	*bilinear = bilinear0;
	return valid;
}

void fill_voxels(int intersection_counter) {

	for (int i = 0; i < intersection_counter; i++) {
		float3 intrs0 = intersections[i*2 + 0];
		float3 intrs1 = intersections[i*2 + 1];
		intrs0 = scale(1/volume_spacing, intrs0);
		intrs1 = scale(1/volume_spacing, intrs1);

		for (int z = intrs0.z; z <= intrs1.z; z++)
			for (int y = intrs0.y; y <= intrs1.y; y++)
				for (int x = intrs0.x; x <= intrs1.x; x++) {
					float3 voxel_coord = {x*volume_spacing,y*volume_spacing,z*volume_spacing};
					if (inrange(x, 0, volume_w) && inrange(y, 0, volume_h) && inrange(z, 0, volume_n)) {
						float dists[dw_n];
						unsigned char bilinears[dw_n];
						bool valid = true;
						float G = 0;
						unsigned char contribution = 0;
						for (int n = 0; n < dw_n; n++) {
							valid &= project(&bilinears[n], &dists[n], voxel_coord, n);
							G += 1/dists[n];
						}
						if (!valid) continue;

						for (int a = 0; a < dw_n; a++) contribution += 1/dists[a]/G*bilinears[a];
						
						//if (volume_a(x,y,z) != 0) volume_a(x,y,z) = (volume_a(x,y,z) + contribution)/2;	else volume_a(x,y,z) = contribution;
						//volume_a(x,y,z) = max(volume_a(x,y,z), contribution);
						//if (volume_a(x,y,z) == 0) volume_a(x,y,z) = contribution;
						volume_a(x,y,z) = contribution;
					}
				}
	}
}

void reconstruct() {
	holger_time_start(1, "Reconstruction");

	memset(volume, 0, sizeof(unsigned char)*volume_w*volume_h*volume_n);

	printf("bscan_w/h: %d %d\n", bscan_w, bscan_h);
	printf("volume_w/h/n: %d %d %d\n", volume_w, volume_h, volume_n);
	printf("volume_spacing: %f\n", volume_spacing);
	printf("volume_origo: %f %f %f\n", volume_origo.x, volume_origo.y, volume_origo.z);
	printf("\n");

	holger_time(1, "Initialization");

	// Fill up queue
	for (int i = 0; i < dw_n; i++) {
		shift_queues();
		wait_for_input();
		if (end_of_data()) break;
	}

	while(true) {
		#ifdef PRINT 
		printf("Reconstructing\n"); 
		#endif
		
		// Retrieve ultrasound data and perform ye olde switcheroo
		shift_queues();
		wait_for_input();
		if (end_of_data()) break;

		holger_time(1, "Retrieve ultrasound data");

		calibrate_pos_matrix(pos_matrices_queue[2], cal_matrix);
		holger_time(1, "Calibrate");

		create_plane_points(pos_matrices_queue[2]);
		holger_time(1, "Fill, transform and translate plane_points");

		create_plane_eq();
		holger_time(1, "Fill bscan_plane_equation");

		//int axis = find_orthogonal_axis(bscan_plane_equation); // TODO
		int axis = 1;

		int intersection_counter = find_intersections(axis);
		holger_time(1, "Find ray-plane intersections");

		fill_voxels(intersection_counter);
		holger_time(1, "Fill voxels");

		#ifdef PRINT 
		printf("\n");
		#endif
	}

	holger_time_print(1);
}